{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import all libraries\n",
    "import scispacy\n",
    "import spacy\n",
    "import json\n",
    "import pandas as pd\n",
    "from nltk.stem.snowball import SnowballStemmer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the sciSpacy model\n",
    "nlp = spacy.load(\"en_core_sci_sm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load all articles from the specified category\n",
    "articles = []\n",
    "category = 'nucl-ex'\n",
    "with open(\"arxiv-metadata-oai-snapshot.json\", \"r\") as f:\n",
    "    for l in f:\n",
    "        d = json.loads(l)\n",
    "        if category in d['categories'].split(' '):\n",
    "            articles.append(d)\n",
    "df = pd.DataFrame.from_records(articles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean the abstracts and prepare for Spacy pipe\n",
    "df['clean_abstract'] = [x.replace('\\n',' ').strip() for x in df['abstract']]\n",
    "df['tuple_input'] = [(row['clean_abstract'], row['id']) for _,row in df.iterrows()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SciSpacy NER extraction\n",
    "results = []\n",
    "for doc, context in nlp.pipe(df['tuple_input'].to_list(), as_tuples=True):\n",
    "    results.append({'id':context, 'entities':doc.ents})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "stemmer = SnowballStemmer(\"english\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean the results by checking if token is longer than 1 character and consists of alphabetic characters\n",
    "# Stem the filtered tokens\n",
    "exploded = []\n",
    "for row in results:\n",
    "    for el in row['entities']:\n",
    "        if el.text.isalpha() and len(el.text) > 1:\n",
    "            exploded.append({'id':row['id'], 'entity':\" \".join([stemmer.stem(word) for word in el.text.split(' ')])})\n",
    "results_df = pd.DataFrame.from_records(exploded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>entity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0704.0075</td>\n",
       "      <td>sector</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0704.0075</td>\n",
       "      <td>heavi</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0704.0075</td>\n",
       "      <td>baryon</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0704.0075</td>\n",
       "      <td>year</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0704.0075</td>\n",
       "      <td>decay</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          id  entity\n",
       "0  0704.0075  sector\n",
       "1  0704.0075   heavi\n",
       "2  0704.0075  baryon\n",
       "3  0704.0075    year\n",
       "4  0704.0075   decay"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>unique_values</th>\n",
       "      <th>counts</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>measur</td>\n",
       "      <td>13382</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>data</td>\n",
       "      <td>9518</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>collis</td>\n",
       "      <td>8215</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>result</td>\n",
       "      <td>6153</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>experi</td>\n",
       "      <td>5909</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>product</td>\n",
       "      <td>5622</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>experiment</td>\n",
       "      <td>4962</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>energi</td>\n",
       "      <td>4373</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>calcul</td>\n",
       "      <td>4129</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>nuclei</td>\n",
       "      <td>4060</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>interact</td>\n",
       "      <td>3829</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>studi</td>\n",
       "      <td>3730</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>effect</td>\n",
       "      <td>3557</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>model</td>\n",
       "      <td>3533</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>correl</td>\n",
       "      <td>3279</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>neutron</td>\n",
       "      <td>3151</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>distribut</td>\n",
       "      <td>3106</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>gev</td>\n",
       "      <td>3105</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>state</td>\n",
       "      <td>3044</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>investig</td>\n",
       "      <td>2896</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>rhic</td>\n",
       "      <td>2722</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>particl</td>\n",
       "      <td>2703</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>central</td>\n",
       "      <td>2545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>paramet</td>\n",
       "      <td>2523</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>detector</td>\n",
       "      <td>2511</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>reaction</td>\n",
       "      <td>2464</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>proton</td>\n",
       "      <td>2454</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>function</td>\n",
       "      <td>2447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>lhc</td>\n",
       "      <td>2329</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td>predict</td>\n",
       "      <td>2322</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   unique_values  counts\n",
       "0         measur   13382\n",
       "1           data    9518\n",
       "2         collis    8215\n",
       "3         result    6153\n",
       "4         experi    5909\n",
       "5        product    5622\n",
       "6     experiment    4962\n",
       "7         energi    4373\n",
       "8         calcul    4129\n",
       "9         nuclei    4060\n",
       "10      interact    3829\n",
       "11         studi    3730\n",
       "12        effect    3557\n",
       "13         model    3533\n",
       "14        correl    3279\n",
       "15       neutron    3151\n",
       "16     distribut    3106\n",
       "17           gev    3105\n",
       "18         state    3044\n",
       "19      investig    2896\n",
       "20          rhic    2722\n",
       "21       particl    2703\n",
       "22       central    2545\n",
       "23       paramet    2523\n",
       "24      detector    2511\n",
       "25      reaction    2464\n",
       "26        proton    2454\n",
       "27      function    2447\n",
       "28           lhc    2329\n",
       "29       predict    2322"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the count of entities\n",
    "counts = results_df['entity'].value_counts().rename_axis('unique_values').reset_index(name='counts')\n",
    "counts.head(30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove words that appear fewer than 5 times or more than 4500\n",
    "remove_words = counts[(counts['counts'] < 5) | (counts['counts'] > 4500)]['unique_values'].to_list()\n",
    "final_df = results_df.drop(results_df[results_df['entity'].isin(remove_words)].index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "energi       4373\n",
       "calcul       4129\n",
       "nuclei       4060\n",
       "interact     3829\n",
       "studi        3730\n",
       "             ... \n",
       "crustal         5\n",
       "scd             5\n",
       "jedi            5\n",
       "lbne            5\n",
       "inventori       5\n",
       "Name: entity, Length: 3695, dtype: int64"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check the results\n",
    "final_df['entity'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_df.to_csv('nucl-ex-entities.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "scispacy",
   "language": "python",
   "name": "scispacy"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
